from typing import Any, Callable, Iterable, Mapping

from zkl_pyutils_fsspec import FsLike, resolve_fs

from zkl_aiutils_datasets.utils.qual import get_and_check_module_and_qualname
from .loading import dataset_script_file_name


def write_dataset_script(fs: FsLike, dataset_script: str):
    fs = resolve_fs(fs)
    fs.mkdirs("", exist_ok=True)
    with fs.open(dataset_script_file_name, "wt") as fp:
        fp.write(dataset_script)


def make_dataset_script(
    loader: Callable | type,
    args: Iterable[Any] | None = None,
    kwargs: Mapping[str, Any] | None = None,
) -> str:
    module, qualname = get_and_check_module_and_qualname(loader)

    lines = []
    lines.append(f"import {module} as loader_module")
    lines.append("")
    lines.append("def load(fs, *args, **kwargs):")
    if args is not None:
        args_code = ', '.join(repr(arg) for arg in args)
        lines.append(f"    args = ({args_code}, *args)")
    if kwargs is not None:
        kwargs_code = ', '.join(f'{repr(key)}: {repr(arg)}' for key, arg in kwargs.items())
        lines.append(f"    kwargs = {{{kwargs_code}, **kwargs}}")
    lines.append(f"    loader = loader_module.{qualname}")
    lines.append(f"    return loader(fs, *args, **kwargs)")
    lines.append("")
    return "\n".join(lines)
